#!/usr/bin/python

import collections

class QuickUnionWeighted(object):

  def __init__(self, num_objects):
    self.num_objects = num_objects
    self.id_list = range(self.num_objects)
    self.id_to_size = {}
    for obj_id in self.id_list:
      self.id_to_size[obj_id] = 1

  def GetRoot(self, object_id):
    while not self.id_list[object_id] == object_id:
      object_id = self.id_list[object_id]
    return object_id

  def connected(self, p, q):
    return self.GetRoot(p) == self.GetRoot(q)

  def union(self, p, q):
    print 'Connecting %d and %d' % (p, q)
    p_root = self.GetRoot(p)
    q_root = self.GetRoot(q)
    p_size = self.id_to_size[p_root]
    q_size = self.id_to_size[q_root]
    if p_size > q_size:
      self.id_list[q_root] = p_root
      self.id_to_size[p_root] += q_size
    else:
      self.id_list[p_root] = q_root
      self.id_to_size[q_root] += p_size
    print 'pos: ', range(self.num_objects)
    print 'ids: ', self.id_list
    print 'sizes: ', self.id_to_size

def main():
  num_objects = 10
  quw = QuickUnionWeighted(num_objects)
  quw.union(1, 3)
  quw.union(2, 4)
  quw.union(1, 4)
  quw.union(1, 9)
  print '1, 9 connected: ', quw.connected(1, 9)
  print '1, 4 connected: ', quw.connected(1, 4)

if __name__ == '__main__':
  main() 
